// Bandwidth control and rate limiting

use std::time::{Duration, Instant};
use tokio::time::sleep;

pub struct BandwidthLimiter {
    rate: u64,  // bytes per second
    last_update: Instant,
    tokens: f64,
}

impl BandwidthLimiter {
    pub fn new(rate: u64) -> Self {
        Self {
            rate,
            last_update: Instant::now(),
            tokens: rate as f64,
        }
    }

    pub async fn acquire(&mut self, bytes: usize) {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_update).as_secs_f64();
        
        // Refill tokens
        self.tokens += elapsed * self.rate as f64;
        self.tokens = self.tokens.min(self.rate as f64);
        self.last_update = now;

        // Wait if not enough tokens
        if (bytes as f64) > self.tokens {
            let wait_time = ((bytes as f64 - self.tokens) / self.rate as f64) * 1000.0;
            sleep(Duration::from_millis(wait_time as u64)).await;
            self.tokens = 0.0;
        } else {
            self.tokens -= bytes as f64;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_bandwidth_limiter() {
        let mut limiter = BandwidthLimiter::new(1000); // 1KB/s
        
        let start = Instant::now();
        limiter.acquire(500).await;
        limiter.acquire(500).await;
        let elapsed = start.elapsed();
        
        // Should take about 1 second
        assert!(elapsed.as_millis() >= 900);
    }
}
